#include "tcp_transfer.h"

// 服务器（接收文件）
void tcp_server() {
    int server_fd, client_fd;
    struct sockaddr_in server_addr, client_addr;
    socklen_t client_len = sizeof(client_addr);

    // 创建套接字
    server_fd = socket(AF_INET, SOCK_STREAM, 0);
    if (server_fd == -1) {
        perror("创建套接字失败");
        return;
    }

    // 端口复用
    int opt = 1;
    setsockopt(server_fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));

    // 绑定地址
    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_addr.s_addr = INADDR_ANY;
    server_addr.sin_port = htons(TCP_PORT);

    if (bind(server_fd, (struct sockaddr*)&server_addr, sizeof(server_addr)) == -1) {
        perror("绑定失败");
        close(server_fd);
        return;
    }

    // 监听并接受连接
    if (listen(server_fd, 5) == -1) {
        perror("监听失败");
        close(server_fd);
        return;
    }
    printf("服务器启动，监听端口 %d...\n", TCP_PORT);
    client_fd = accept(server_fd, (struct sockaddr*)&client_addr, &client_len);
    if (client_fd == -1) {
        perror("接受连接失败");
        close(server_fd);
        return;
    }
    printf("客户端 %s:%d 已连接\n",
           inet_ntoa(client_addr.sin_addr), ntohs(client_addr.sin_port));

    // 接收文件
    FileTransferInfo info;
    while (recv(client_fd, &info, sizeof(info), 0) > 0) {
        printf("\n接收文件: %s （大小: %llu bytes）\n",
               info.filename, (unsigned long long)info.file_size);

        int fd = open(info.filename, O_WRONLY | O_CREAT | O_TRUNC, 0644);
        if (fd == -1) {
            perror("创建文件失败");
            break;
        }

        char buff[BUFFER_SIZE];
        uint64_t recv_size = 0;
        ssize_t ret;
        while (recv_size < info.file_size) {
            ret = recv(client_fd, buff, sizeof(buff), 0);
            if (ret <= 0) {
                perror("接收数据失败");
                break;
            }
            write(fd, buff, ret);
            recv_size += ret;
            print_progress((double)recv_size / info.file_size * 100);
        }
        close(fd);
        printf("\n文件接收完成\n");
    }

    close(client_fd);
    close(server_fd);
}

// 客户端（发送文件）
void tcp_client() {
    char server_ip[32], filename[256];
    printf("\n===== 传输客户端 =====");
    printf("\n服务器IP: ");
    scanf("%s", server_ip);
    printf("发送文件路径: ");
    scanf("%s", filename);

    // 校验文件
    if (access(filename, R_OK) == -1) {
        perror("文件不可读");
        return;
    }

    // 获取文件信息
    FileTransferInfo info;
    memset(&info, 0, sizeof(info));
    char* basename = strrchr(filename, '/');
    strcpy(info.filename, basename ? basename + 1 : filename);
    info.file_size = get_file_size(filename);
    if (info.file_size == -1) return;

    // 连接服务器
    int sock = socket(AF_INET, SOCK_STREAM, 0);
    if (sock == -1) {
        perror("创建套接字失败");
        return;
    }

    struct sockaddr_in server_addr;
    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(TCP_PORT);
    if (inet_pton(AF_INET, server_ip, &server_addr.sin_addr) <= 0) {
        perror("无效IP");
        close(sock);
        return;
    }

    if (connect(sock, (struct sockaddr*)&server_addr, sizeof(server_addr)) == -1) {
        perror("连接失败");
        close(sock);
        return;
    }

    // 发送文件
    send(sock, &info, sizeof(info), 0);
    int fd = open(filename, O_RDONLY);
    if (fd == -1) {
        perror("打开文件失败");
        close(sock);
        return;
    }

    char buff[BUFFER_SIZE];
    uint64_t sent_size = 0;
    ssize_t ret;
    while ((ret = read(fd, buff, sizeof(buff))) > 0) {
        send(sock, buff, ret, 0);
        sent_size += ret;
        print_progress((double)sent_size / info.file_size * 100);
    }
    close(fd);
    close(sock);
    printf("\n文件发送完成\n");
}

// 传输模块入口
void tcp_transfer_menu() {
    int choice;
    printf("\n===== 局域网传输 =====");
    printf("\n1. 启动服务器（接收）");
    printf("\n2. 启动客户端（发送）");
    printf("\n3. 返回主菜单");
    printf("\n选择: ");

    if (scanf("%d", &choice) != 1) {
        choice = -1;
        while (getchar() != '\n');
    }

    switch (choice) {
        case 1: tcp_server(); break;
        case 2: tcp_client(); break;
        case 3: return;
        default: printf("无效选择\n");
    }
}
